#include <iostream>
#include <cstring>
#include <string.h>
#include <unistd.h>
#include <chrono>
#include <fcntl.h>
#include <arpa/inet.h>
#include "tzsptap2pcap.hpp"
#include "common.hpp"
#include "config.h"
#include "log.h"
#include "json/json.h"

using namespace std;

// 初始化 PCAP 文件头
void init_pcap_file_header(pcap_file_header& header)
{
	header.magic_number = 0xa1b2c3d4;
	header.version_major = 2;
	header.version_minor = 4;
	header.thiszone = 0;
	header.sigfigs = 0;
	header.snaplen = 65535;
	header.network = 1;
}

int parse_tzsppkg(char *data, size_t datalen, tzsppkg_t *pkg)
{
	if (!data)
		return -1;
	if (datalen < 4)
		return -1;
	if (!pkg)
		return -1;

	pkg->num_fields = 0;
	pkg->tzsphdr = (tzspheader *) data;
	data += 4;
	datalen -= 4;

	if (!(pkg->tzsphdr->flags & TZSP_FLAG_NOFIELDS)) {
		for (int n = 0; n < TZSP_MAX_FIELDS; n++) {
			pkg->num_fields += 1;
			if (datalen < 1)
				return -1;
			pkg->fields[n] = (struct tzspfield *) data;
			if (pkg->fields[n]->tag > 2) { //not PAD nor END; carries len and data
				data += (1 + 1 + pkg->fields[n]->len);
				datalen -= (1 + 1 + pkg->fields[n]->len);
			} else {
				data += 1; datalen -= 1;
				if (pkg->fields[n]->tag == TZSP_HDR_END)
					break;
			}
		}
	}

	if (datalen > 0) {
		pkg->data = data;
		pkg->datalen = datalen;
	} else {
		pkg->data = nullptr;
		pkg->datalen = 0;
	}

	return 0;
}
tzsptap2pcap::tzsptap2pcap(/* args */)
{
	pcap_file_size = 0;
	pcap_writting = nullptr;
	pcap_new_write = nullptr;
	stopThreadFlag.store(false, std::memory_order_seq_cst);
	sys_debug("New tzsptap2pcap default!\n");
	Lzlog_debug("New tzsptap2pcap default!");
}

tzsptap2pcap::tzsptap2pcap(string &dev_ip, char* &pcap_sigle_size_str)
{
	pcap_writting = nullptr;
	pcap_new_write = nullptr;
	stopThreadFlag.store(false, std::memory_order_seq_cst);
	tzsp_dev_ip = dev_ip;
	pcap_file_size = BytesStringToBytes(pcap_sigle_size_str);
	sys_debug("New tzsptap2pcap set ip[%s], pcap_single_size[%s]!\n",
		tzsp_dev_ip.c_str(), pcap_sigle_size_str);
	Lzlog_debug("New tzsptap2pcap set ip[%s], pcap_single_size[%s]!",
		tzsp_dev_ip.c_str(), pcap_sigle_size_str);
}

tzsptap2pcap::~tzsptap2pcap()
{
	Lzlog_debug("Deleting tzsptap2pcap_device[ip:%s]", tzsp_dev_ip.c_str());

	stop_thread();
	srcpkg_list_cv.notify_all();
	sys_debug("Clean srcpkg_list;\n");
	/* 清空 srcpkg_list */
	for (auto it = srcpkg_list.begin(); it != srcpkg_list.end(); it++) {
		delete *it;
	}
	srcpkg_list.clear();

        for (auto it = pcap_files.begin(); it != pcap_files.end(); it++) {
                delete *it;
        }
        pcap_files.clear();

	sys_debug("Close pacp outFile.\n ");
	sys_debug("Deleted tzsptap2pcap_device[ip:%s]:\n", tzsp_dev_ip.c_str());
}

void tzsptap2pcap::set_pcap_save_top_dir(const char *dir_path)
{
	pcap_save_top_dir = dir_path;
	if (!pcap_save_top_dir.empty() && pcap_save_top_dir.back() != '/') {
		pcap_save_top_dir.push_back('/');
	}
        if (pkg_flag == PKG_TYPE_TZSP) {
                pcap_save_top_dir += "TZSP/";
        } else if (pkg_flag == PKG_TYPE_SHARK) {
               pcap_save_top_dir += "PCAP/";
        }
	pcap_save_top_dir.append(tzsp_dev_ip);
	create_directory(tzsp_dev_ip.c_str());
}
void tzsptap2pcap::compose_open_file(timeval _ts)
{
	time_t time_secend = _ts.tv_sec;
	tm *tm_info = localtime(&time_secend);

	char file_buf[20] = {0};
	strftime(file_buf, sizeof(file_buf), "/%02H-%02M-%02S.pcap", tm_info);
	string pcap_dts_file_path = pcap_save_back_dir + string(file_buf);

	char tmp_file_buf[20] = {0};
	strftime(tmp_file_buf, sizeof(tmp_file_buf), "/.%02H-%02M-%02S.pcap", tm_info);
	string pcap_dts_tmp_file_path = pcap_save_back_dir + string(tmp_file_buf);
	create_directory(pcap_save_back_dir.c_str());
	pcap_file_save_t *pcap_file = new pcap_file_save_t(pcap_dts_tmp_file_path, pcap_dts_file_path);
	if (!pcap_file)
		return;
	if (!pcap_file->outFile.is_open()) {
		delete pcap_file;
	}
	pcap_file_header file_header;
	init_pcap_file_header(file_header);
	pcap_file->outFile.write(reinterpret_cast<const char*>(&file_header), sizeof(file_header));
	if (pcap_file->outFile.fail()) {
		sys_err("Write pcap file header to file(%s) failed.\n", pcap_file->file_path.c_str());
		Lzlog_error("Write pcap file header to file(%s) failed.", pcap_file->file_path.c_str());
		delete pcap_file;
		return;
	}

	sys_debug("Creat and open new pcap file(%s) to save pkgs.\n", pcap_dts_file_path.c_str());
	Lzlog_debug("Creat and open new pcap file(%s) to save pkgs.", pcap_dts_file_path.c_str());
	pcap_new_write = pcap_file;
        pcap_files.push_back(pcap_new_write); //添加新的文件到list中
}

/* 将src pkg 中的数据直接写入文件 */
void tzsptap2pcap::write_srcpkg_t_to_pcap_file(srcpkg_t &pkg)
{
	unique_lock<mutex> lock(pcap_writting_mtx);
	if (!pcap_writting) {
                if (!pcap_new_write) {
                        compose_open_file(static_cast<timeval>(pkg.current_time));
                }
	}
	if (pcap_new_write) {
		if (pcap_writting) {
                        pcap_writting->close_rename_file(); //关闭文件描述符，并重命名文件
			//delete pcap_writting;
		}
		pcap_writting = pcap_new_write;
		pcap_new_write = nullptr;
	}
	if (!pcap_writting || !pcap_writting->outFile.is_open()) {
		sys_debug("Not found pcap file is opened, Nothing to do!\n");
		Lzlog_debug("Not found pcap file is opened, Nothing to do!");
		return;
	}
	sys_debug("write src data to PCAP file.(size:%lu Bytes)\n", pkg.datalen);
	Lzlog_debug("write src data to PCAP file.(size:%lu Bytes)", pkg.datalen);
	//dump_hex((const uint8_t *)pkg.data, pkg.datalen, 16);
	printf("\n");

	streampos currentPos = pcap_writting->outFile.tellp();
	ssize_t write_size = 0;

	// 写入 PCAP 数据包头部
	pcap_pkthdr pcap_packet_header;
	pcap_packet_header.time_sec = static_cast<uint32_t>(pkg.current_time.tv_sec);
	pcap_packet_header.time_usec = static_cast<uint32_t>(pkg.current_time.tv_usec);
	pcap_packet_header.caplen = static_cast<uint32_t>(pkg.datalen);
	pcap_packet_header.len = static_cast<uint32_t>(pkg.datalen);
	pcap_writting->outFile.write(reinterpret_cast<const char*>(&pcap_packet_header), sizeof(pcap_pkthdr));
	if (pcap_writting->outFile.fail()) {
		sys_err("Write pcap header to file %s failed, seek fd to before site.\n", pcap_writting->file_path.c_str());
		Lzlog_error("Write pcap header to file %s failed, seek fd to before site.", pcap_writting->file_path.c_str());
		pcap_writting->outFile.seekp(currentPos, ios::beg);
		return;
	}
	write_size += sizeof(pcap_pkthdr);

	// 写入数据
	pcap_writting->outFile.write(reinterpret_cast<const char*>(pkg.data), pkg.datalen);
	if (pcap_writting->outFile.fail()) {
		sys_err("Write pkg payload to file %s failed, seek fd to before site.\n", pcap_writting->file_path.c_str());
		Lzlog_error("Write pkg payload to file %s failed, seek fd to before site.", pcap_writting->file_path.c_str());
		pcap_writting->outFile.seekp(currentPos, ios::beg);
		return;
	}
	write_size += pkg.datalen;
	pcap_writting->count++;
	pcap_writting->size += write_size;

	if (pcap_writting->size >= pcap_file_size) {
		sys_debug("%lu pkg saved in file %s size(%lu B), exceed single size (%lu B), Need close and open new file!\n",
			pcap_writting->count, pcap_writting->file_path.c_str(), pcap_writting->size, pcap_file_size);
		Lzlog_debug("%lu pkg saved in file %s size(%lu B), exceed single size (%lu B), Need close and open new file!",
			pcap_writting->count, pcap_writting->file_path.c_str(), pcap_writting->size, pcap_file_size);
		delete pcap_writting;
		pcap_writting = nullptr;
	} else {
		sys_info("Have %lu pkg saved in file %s size(%lu B)\n", pcap_writting->count, pcap_writting->file_path.c_str(), pcap_writting->size);
		Lzlog_info("Have %lu pkg saved in file %s size(%lu B)", pcap_writting->count, pcap_writting->file_path.c_str(), pcap_writting->size);
	}
	return;
}

/* 将TZSP数据写入到pcap文件中 */
void tzsptap2pcap::write_tzsp_to_pcap_file(tzsppkg_t &pkg)
{
	unique_lock<mutex> lock(pcap_writting_mtx);
	if (!pcap_writting) {
		compose_open_file(static_cast<timeval>(pkg.current_time));
	}
	if (pcap_new_write) {
		if (pcap_writting) {
                        pcap_writting->close_rename_file(); //关闭文件描述符，并重命名文件
			//delete pcap_writting;
		}
		pcap_writting = pcap_new_write;
		pcap_new_write = nullptr;
	}
	if (!pcap_writting || !pcap_writting->outFile.is_open()) {
		sys_debug("Not found pcap file is opened, Nothing to do!\n");
		Lzlog_debug("Not found pcap file is opened, Nothing to do!");
		return;
	}
	sys_debug("Data after TZSP parse write to PCAP file.(size:%lu Bytes)\n", pkg.datalen);
	Lzlog_debug("Data after TZSP parse write to PCAP file.(size:%lu Bytes)", pkg.datalen);
	dump_hex((const uint8_t *)pkg.data, pkg.datalen, 16);
	printf("\n");

	streampos currentPos = pcap_writting->outFile.tellp();
	ssize_t write_size = 0;

	// 写入 PCAP 数据包头部
	pcap_pkthdr pcap_packet_header;
	pcap_packet_header.time_sec = static_cast<uint32_t>(pkg.current_time.tv_sec);
	pcap_packet_header.time_usec = static_cast<uint32_t>(pkg.current_time.tv_usec);
	pcap_packet_header.caplen = static_cast<uint32_t>(pkg.datalen);
	pcap_packet_header.len = static_cast<uint32_t>(pkg.datalen);
	pcap_writting->outFile.write(reinterpret_cast<const char*>(&pcap_packet_header), sizeof(pcap_pkthdr));
	if (pcap_writting->outFile.fail()) {
		sys_err("Write pcap header to file %s failed, seek fd to before site.\n", pcap_writting->file_path.c_str());
		Lzlog_error("Write pcap header to file %s failed, seek fd to before site.", pcap_writting->file_path.c_str());
		pcap_writting->outFile.seekp(currentPos, ios::beg);
		return;
	}
	write_size += sizeof(pcap_pkthdr);
	// 写入数据
	pcap_writting->outFile.write(reinterpret_cast<const char*>(pkg.data), pkg.datalen);
	if (pcap_writting->outFile.fail()) {
		sys_err("Write pkg payload to file %s failed, seek fd to before site.\n", pcap_writting->file_path.c_str());
		Lzlog_error("Write pkg payload to file %s failed, seek fd to before site.", pcap_writting->file_path.c_str());
		pcap_writting->outFile.seekp(currentPos, ios::beg);
		return;
	}
	write_size += pkg.datalen;
	pcap_writting->count++;
	pcap_writting->size += write_size;
	if (pcap_writting->size >= pcap_file_size) {
		sys_debug("%lu pkg saved in file %s size(%lu B), exceed single size (%lu B), Need close and open new file!\n",
			pcap_writting->count, pcap_writting->file_path.c_str(), pcap_writting->size, pcap_file_size);
		Lzlog_debug("%lu pkg saved in file %s size(%lu B), exceed single size (%lu B), Need close and open new file!",
			pcap_writting->count, pcap_writting->file_path.c_str(), pcap_writting->size, pcap_file_size);
		delete pcap_writting;
		pcap_writting = nullptr;
	} else {
		sys_info("Have %lu pkg saved in file %s size(%lu B)\n", pcap_writting->count, pcap_writting->file_path.c_str(), pcap_writting->size);
		Lzlog_info("Have %lu pkg saved in file %s size(%lu B)", pcap_writting->count, pcap_writting->file_path.c_str(), pcap_writting->size);
	}
	return;
}

// 向srcpkg_list 末尾添加元素
void tzsptap2pcap::srcpkg_list_push_back(uint8_t *_data, ssize_t _len, timeval c_time = {0, 0}) //向list中添加元素
{
	srcpkg_t *srcpkg = new srcpkg_t(_len);
	if (!srcpkg) {
		return;
	}
	if (!srcpkg->data) {
		delete srcpkg;
		return;
	}
	memcpy(srcpkg->data, _data, _len);
	if (c_time.tv_sec == 0) {
		gettimeofday(&srcpkg->current_time, nullptr);
	} else {
		srcpkg->current_time = c_time;
	}
	unique_lock<mutex> lock(srcpkg_list_mtx);
	srcpkg_list.push_back(srcpkg);
	srcpkg_list_count++;
	srcpkg_list_cv.notify_one();
}

/* 线程 循环检测tzsp数据列表并转换成pcap格式保存文件 */
void tzsptap2pcap::thread_pare_tzsp_to_pcap()
{
	int ret = -1;
	while (true) {
		unique_lock<mutex> lock(srcpkg_list_mtx);
		// 等待任务队列非空或者停止标志置位
		srcpkg_list_cv.wait(lock, [this] { return!srcpkg_list.empty() || stopThreadFlag.load(std::memory_order_seq_cst);});
		if (!srcpkg_list.empty()) {
			srcpkg_t *srcpkg = srcpkg_list.front();
			srcpkg_list.pop_front();
			srcpkg_list_count--;
			tzsppkg_t pkg;
			ret = parse_tzsppkg((char *)srcpkg->data, srcpkg->datalen, &pkg);//解析TSZP数据
			if (!ret) {
				sys_debug("Get and parse TZSP package success, handle to pcap and save.\n");
				Lzlog_debug("Get and parse TZSP package success, handle to pcap and save.");
				pkg.current_time = srcpkg->current_time;
				write_tzsp_to_pcap_file(pkg); //写入到pcap文件
			} else {
				sys_debug("Get and parse TZSP package failed. drop it\n");
				Lzlog_debug("Get and parse TZSP package failed. drop it");
			}
			delete srcpkg;
		} else if (stopThreadFlag.load(std::memory_order_seq_cst)) {
			break;
		}
	}
}
/* 线程 循环检测数据列表并保存文件 */
void tzsptap2pcap::thread_no_pare_to_pcap()
{
	int ret = -1;
	while (true) {
		unique_lock<mutex> lock(srcpkg_list_mtx);
		// 等待任务队列非空或者停止标志置位
		srcpkg_list_cv.wait(lock, [this] { return!srcpkg_list.empty() || stopThreadFlag.load(std::memory_order_seq_cst);});
		if (!srcpkg_list.empty()) {
			srcpkg_t *srcpkg = srcpkg_list.front();
			srcpkg_list.pop_front();
			srcpkg_list_count--;
                        if (is_tcpdump_data((const char *)srcpkg->data, srcpkg->datalen)) {
			        write_srcpkg_t_to_pcap_file(*srcpkg);
                        } else {
                               sys_debug("srcpkg->data %s format is incorrect, Not add to file\n", srcpkg->data);
                        }
			delete srcpkg;
		} else if (stopThreadFlag.load(std::memory_order_seq_cst)) {
			break;
		}
	}
}

/* 线程 循环检测时间并更新目录 */
void tzsptap2pcap::thread_check_time_up_directory()
{
	int lastHour = -1;
	int current_hour;
	while (true) {
		if (stopThreadFlag.load(std::memory_order_seq_cst)) {
			return;
		}
		auto now = chrono::system_clock::now();
		time_t now_c = chrono::system_clock::to_time_t(now);
		tm* local_time = localtime(&now_c);
		current_hour = local_time->tm_hour;
		if (current_hour != lastHour) {
			char buffer[20] = {0};
			strftime(buffer, sizeof(buffer), "/%Y-%02m-%02d-%02H", local_time);
			string folderName = string(buffer);
			pcap_save_back_dir = pcap_save_top_dir + folderName;
			//create_directory(pcap_save_back_dir.c_str());
			lastHour = current_hour;
			if (pcap_writting) {
				pcap_file_save_t* temp = nullptr;
				{
					unique_lock<mutex> lock(pcap_writting_mtx);
					temp = pcap_writting;
					pcap_writting = nullptr;
				}
				delete temp;
			}
			sys_debug("Creat new directory(%s) to save pcap file.\n", pcap_save_back_dir.c_str());
			Lzlog_debug("Creat new directory(%s) to save pcap file.", pcap_save_back_dir.c_str());
		}
		this_thread::sleep_for(chrono::milliseconds(100));
	}
}

void tzsptap2pcap::stop_thread()
{
	stopThreadFlag.store(true, std::memory_order_seq_cst);

	this_thread::sleep_for(chrono::milliseconds(101));
	/* 等待两个线程 退出并回收 */
	if (t_check_time_up_directory.joinable()) {
		t_check_time_up_directory.join();
	}

	if (t_pare_tzsp_to_pcap.joinable()) {
		t_pare_tzsp_to_pcap.join();
	}
	sys_debug("Stoped Thread check_time_up_directory;\n");
	sys_debug("Stoped Thread pare_tzsp_to_pcap;\n");
}

void tzsptap2pcap::start_thread()
{
	/* 创建线程 检测时间，获取小时并更新目录 /$Top_dir/$dev_IP/$Y-$02m-$02d-$02H */
	t_check_time_up_directory = thread(&tzsptap2pcap::thread_check_time_up_directory, this);
	t_check_time_up_directory.detach();

	if (pkg_flag == PKG_TYPE_TZSP) {
		/* 创建线程 循环检测TZSP src pkg list中是否存在需要处理的数据，有则转换成pcap格式存储到文件中 */
		t_pare_tzsp_to_pcap = thread(&tzsptap2pcap::thread_pare_tzsp_to_pcap, this);
	} else {
		t_pare_tzsp_to_pcap = thread(&tzsptap2pcap::thread_no_pare_to_pcap, this);
	}
	t_pare_tzsp_to_pcap.detach();

	return;
}

void tzsptap2pcap::state()
{
        Json::Value root = pcap_files_formatJSON();
        root["top_dir"] = pcap_save_top_dir;
        if (stopThreadFlag) {
                root["state"] = "2";
        } else {
                root["state"] = "1";
        }
        // 将JSON对象转换为字符串
        Json::StreamWriterBuilder builder;
        builder["commentStyle"] = "None";//关闭注释
        builder["indentation"] = "  ";  // 缩进2 字符
        builder["emitUTF8"] = true;// 使用UTF-8编码
        string jsonString = Json::writeString(builder, root);
        sys_debug("%s\n", jsonString.c_str());
}
/* 将抓包已经写入的文件信息列表组合成JSON */
Json::Value tzsptap2pcap::pcap_files_formatJSON()
{
        Json::Value root;
        uint64_t pkg_count = 0;
        uint64_t pkg_size = 0;
        root["ipaddr"] = tzsp_dev_ip;
        if (pkg_flag == PKG_TYPE_TZSP) {
                root["pkg_type"] = "TZSP";
        } else if (pkg_flag == PKG_TYPE_SHARK) {
                root["pkg_type"] = "SHARK";
        }
        root["files_count"] = to_string(pcap_files.size());
        Json::Value filesarray(Json::arrayValue);
        if (!pcap_files.empty()) {
                for (const auto& pcap_file : pcap_files) {
                        if (pcap_file) {
                                pkg_count += pcap_file->count;
                                pkg_size += pcap_file->size;
                                Json::Value obj;
                                obj["size"] = to_string(pcap_file->size);
                                obj["count"] = to_string(pcap_file->count);
                                obj["path"] = pcap_file->file_path;
                                filesarray.append(obj);
                        }
                }
        }
        root["files"] = filesarray;
        return root;
}

/* 将抓包已经写入的文件信息列表组合成JSON并输入为string类型 */
string tzsptap2pcap::pcap_files_formatJSON_string()
{
        Json::Value root = pcap_files_formatJSON();
        // 将JSON对象转换为字符串
        Json::StreamWriterBuilder builder;
        builder["commentStyle"] = "None";//关闭注释
        builder["indentation"] = "";  // 关闭缩进
        builder["emitUTF8"] = true;// 使用UTF-8编码
        string jsonString = Json::writeString(builder, root);
        sys_debug("获取抓包保存信息JSON数据:%s\n", jsonString.c_str());
        return jsonString;
}